import rclpy
from rclpy.node import Node
from sensor_msgs.msg import Imu
from sensor_msgs.msg import Image,CameraInfo
from sensor_msgs.msg import PointCloud2,PointField
from std_msgs.msg import Header
import sensor_msgs_py.point_cloud2 as pc2
import struct

from rclpy.qos import QoSProfile, QoSDurabilityPolicy, QoSReliabilityPolicy, QoSHistoryPolicy
from rclpy.duration import Duration

class LivoxPointXyzrtlt:
    def __init__(self, x, y, z, reflectivity, tag, line, timestamp):
        self.x = x
        self.y = y
        self.z = z
        self.reflectivity = reflectivity
        self.tag = tag
        self.line = line
        self.timestamp = timestamp

class DataConversion(Node):
    def __init__(self):
        super().__init__('gazebo_data_conversion')

        self.last_stamp = self.get_clock().now().to_msg()

        chen_qos_profile = QoSProfile(
            reliability=QoSReliabilityPolicy.BEST_EFFORT,
            depth=1
        )

        han_qos_profile = QoSProfile(
            history=QoSHistoryPolicy.KEEP_LAST,
            depth=10,
            reliability=QoSReliabilityPolicy.RELIABLE,
            durability=QoSDurabilityPolicy.TRANSIENT_LOCAL,
            deadline=Duration(seconds=1.0)
        )

        self.imu_subscription = self.create_subscription(Imu,'/imu_plugin/out',self.imu_listener_callback,10)
        self.imu_publisher = self.create_publisher(Imu,'/imu_data',han_qos_profile)

        self.subscription = self.create_subscription(PointCloud2,'/points_raw',self.lidar_callback,10)
        self.publisher = self.create_publisher(PointCloud2, '/lidar_point_cloud', han_qos_profile)

        self.rgbd1_depth_image_raw_subscription = self.create_subscription(Image,'/gazebo/rgbd1/depth/image_raw',self.rgbd1_depth_image_raw_listener_callback,10)
        self.rgbd1_depth_image_raw_publisher = self.create_publisher(Image,'/rgbd1/depth/image_rect_raw',chen_qos_profile)

        self.rgbd1_depth_camera_info_subscription = self.create_subscription(CameraInfo,'/gazebo/rgbd1/depth/camera_info',self.rgbd1_depth_camera_info_listener_callback,10)
        self.rgbd1_depth_camera_info_publisher = self.create_publisher(CameraInfo,'/rgbd1/depth/camera_info',chen_qos_profile)

        self.rgbd1_color_image_raw_subscription = self.create_subscription(Image,'/gazebo/rgbd1/color/image_raw',self.rgbd1_color_image_raw_listener_callback,10)
        self.rgbd1_color_image_raw_publisher = self.create_publisher(Image,'/rgbd1/color/image_raw',chen_qos_profile)

        self.rgbd1_color_camera_info_subscription = self.create_subscription(CameraInfo,'/gazebo/rgbd1/color/camera_info',self.rgbd1_color_camera_info_listener_callback,10)
        self.rgbd1_color_camera_info_publisher = self.create_publisher(CameraInfo,'/rgbd1/color/camera_info',chen_qos_profile)

        self.rgbd2_depth_image_raw_subscription = self.create_subscription(Image,'/gazebo/rgbd2/depth/image_raw',self.rgbd2_depth_image_raw_listener_callback,10)
        self.rgbd2_depth_image_raw_publisher = self.create_publisher(Image,'/rgbd2/depth/image_rect_raw',chen_qos_profile)

        self.rgbd2_depth_camera_info_subscription = self.create_subscription(CameraInfo,'/gazebo/rgbd2/depth/camera_info',self.rgbd2_depth_camera_info_listener_callback,10)
        self.rgbd2_depth_camera_info_publisher = self.create_publisher(CameraInfo,'/rgbd2/depth/camera_info',chen_qos_profile)

        self.rgbd2_color_image_raw_subscription = self.create_subscription(Image,'/gazebo/rgbd2/color/image_raw',self.rgbd2_color_image_raw_listener_callback,10)
        self.rgbd2_color_image_raw_publisher = self.create_publisher(Image,'/rgbd2/color/image_raw',chen_qos_profile)

        self.rgbd2_color_camera_info_subscription = self.create_subscription(CameraInfo,'/gazebo/rgbd2/color/camera_info',self.rgbd2_color_camera_info_listener_callback,10)
        self.rgbd2_color_camera_info_publisher = self.create_publisher(CameraInfo,'/rgbd2/color/camera_info',chen_qos_profile)

    def imu_listener_callback(self, msg):
        self.imu_publisher.publish(msg)

    def lidar_callback(self, msg):
        # 将 `PointCloud2` 消息转换为点列表
        points_xyzirt = list(pc2.read_points(msg, skip_nans=True, field_names=("x", "y", "z", "intensity", "ring", "time")))

        # 检查点云是否为空
        if not points_xyzirt:
            self.get_logger().warn('No points received in points_raw')
            return

        # 获取点云消息的起始时间和结束时间
        start_time = msg.header.stamp.sec + msg.header.stamp.nanosec * 1e-9  # 将时间戳转换为秒
        end_time = start_time + (points_xyzirt[-1][5] - points_xyzirt[0][5])  # 计算时间范围

        # 转换点云数据并为每个点分配线性时间戳
        livox_points = self.convert_xyzirt_to_livox(points_xyzirt, start_time, end_time)

        # 构造新的 PointCloud2 消息
        header = Header()
        header.frame_id = msg.header.frame_id

        converted_msg = self.create_pointcloud2_msg(header, livox_points)

        if converted_msg.header.stamp == self.last_stamp:
            self.get_logger().warn("header.stamp is equal to last_stamp")
        else:
            # 发布新的点云消息
            self.publisher.publish(converted_msg)

        self.last_stamp = converted_msg.header.stamp

    def convert_xyzirt_to_livox(self, points_xyzirt, start_time, end_time):
        # 转换 xyzirt 格式到 LivoxPointXyzrtlt，并为每个点生成时间戳
        livox_points = []
        num_points = len(points_xyzirt)

        for i, point in enumerate(points_xyzirt):
            x, y, z, intensity, ring, timestamp = point

            # 计算每个点的时间戳（线性插入时间）
            point_time = start_time + (i / (num_points - 1)) * (end_time - start_time)

            tag = 0  # 默认 tag 为 0，可根据需要修改
            livox_point = LivoxPointXyzrtlt(
                x=x,
                y=y,
                z=z,
                reflectivity=intensity,
                tag=tag,
                line=ring,
                timestamp=point_time
            )
            livox_points.append(livox_point)

        return livox_points

    def create_pointcloud2_msg(self, header, livox_points):
        # 创建 PointCloud2 消息
        fields = [
            PointField(name='x', offset=0, datatype=PointField.FLOAT32, count=1),
            PointField(name='y', offset=4, datatype=PointField.FLOAT32, count=1),
            PointField(name='z', offset=8, datatype=PointField.FLOAT32, count=1),
            PointField(name='reflectivity', offset=12, datatype=PointField.FLOAT32, count=1),
            PointField(name='tag', offset=16, datatype=PointField.UINT8, count=1),
            PointField(name='line', offset=17, datatype=PointField.UINT8, count=1),
            PointField(name='timestamp', offset=18, datatype=PointField.FLOAT64, count=1),
        ]

        # 打包点云数据
        point_cloud_data = []
        for point in livox_points:
            point_cloud_data.append([point.x, point.y, point.z, point.reflectivity, point.tag, point.line, point.timestamp])

        # 使用最后一个点的时间戳作为点云消息的时间戳
        header.stamp = rclpy.time.Time(seconds=livox_points[-1].timestamp).to_msg()

        # 构建 PointCloud2 消息
        pointcloud_msg = pc2.create_cloud(header, fields, point_cloud_data)
        return pointcloud_msg


    def rgbd1_depth_image_raw_listener_callback(self, msg):
        self.rgbd1_depth_image_raw_publisher.publish(msg)

    def rgbd1_depth_camera_info_listener_callback(self, msg):
        self.rgbd1_depth_camera_info_publisher.publish(msg)

    def rgbd1_color_image_raw_listener_callback(self, msg):
        self.rgbd1_color_image_raw_publisher.publish(msg)

    def rgbd1_color_camera_info_listener_callback(self, msg):
        self.rgbd1_color_camera_info_publisher.publish(msg)
    
    def rgbd2_depth_image_raw_listener_callback(self, msg):
        self.rgbd2_depth_image_raw_publisher.publish(msg)

    def rgbd2_depth_camera_info_listener_callback(self, msg):
        self.rgbd2_depth_camera_info_publisher.publish(msg)

    def rgbd2_color_image_raw_listener_callback(self, msg):
        self.rgbd2_color_image_raw_publisher.publish(msg)

    def rgbd2_color_camera_info_listener_callback(self, msg):
        self.rgbd2_color_camera_info_publisher.publish(msg)
        
def main(args=None):
    rclpy.init(args=args)
    data_conversion = DataConversion()
    rclpy.spin(data_conversion)
    relay_node.destroy_node()
    rclpy.shutdown()

if __name__ == '__main__':
    main()
